from dataclasses import dataclass, field
from typing import List
from dacite import from_dict

@dataclass
class Entity:
    paper_id: str
    paper_title: str
    chunk_id: int
    chunk_text: str
    original_filename: str

@dataclass
class QueryResultItem:
    id: int
    distance: float
    entity: Entity

@dataclass
class QueryResult:
    results: List[QueryResultItem] = field(default_factory=list)

def wrap_search_results(search_results):
    wrapped_search_result = QueryResult(results=[
        from_dict(data_class=QueryResultItem, data=item) 
        for item in search_results
    ])
    return wrapped_search_result

def display_query_results(query_result: QueryResult, 
                        max_items: int = None,
                        show_distance: bool = True,
                        show_chunk_text: bool = True,
                        text_truncate: int = 100) -> str:
    """
    展示查询结果
    
    参数:
        query_result: QueryResult对象
        max_items: 最多显示的结果数量
        show_distance: 是否显示距离值
        show_chunk_text: 是否显示文本块内容
        text_truncate: 文本截断长度
        
    返回:
        str: 格式化的结果字符串
    """
    
    output = []  # 用列表存储所有输出行
    
    results = query_result.results[:max_items] if max_items else query_result.results
    
    if not results:
        return "没有找到相关结果"
        
    output.append(f"\n找到 {len(results)} 条结果：")
    output.append("-" * 80)
    
    for idx, item in enumerate(results, 1):
        # 标题行
        output.append(f"\n[{idx}] {item.entity.paper_title}")
        
        # 基础信息行
        info_parts = [
            f"Paper ID: {item.entity.paper_id}",
            f"Chunk ID: {item.entity.chunk_id}"
        ]
        if show_distance:
            info_parts.append(f"Distance: {item.distance:.4f}")
        output.append(" | ".join(info_parts))
        
        # 文件名
        output.append(f"File: {item.entity.original_filename}")
        
        # 文本内容（如果需要显示）
        if show_chunk_text:
            text = item.entity.chunk_text
            if text_truncate and len(text) > text_truncate:
                text = text[:text_truncate] + "..."
            output.append("\nText:")
            output.append(text)
            
        output.append("-" * 80)
    
    result = "\n".join(output)
    print(result)  # 仍然打印结果
    return result  # 同时返回结果字符串


# 使用示例：
if __name__ == "__main__":
    # 创建示例数据
    sample_data = QueryResult(results=[
        QueryResultItem(
            id=1,
            distance=0.123,
            entity=Entity(
                paper_id="paper001",
                paper_title="Sample Paper Title",
                chunk_id=1,
                chunk_text="This is a sample chunk of text that could be quite long...",
                original_filename="sample.pdf"
            )
        ),
        # 可以添加更多示例项...
    ])
    
    # 展示结果
    result = display_query_results(
        sample_data,
        max_items=5,
        show_distance=True,
        show_chunk_text=True,
        text_truncate=100
    )

# # 使用示例
# data = [{
#     "id": 1,
#     "distance": 0.75,
#     "entity": {
#         "paper_id": "paper123",
#         "paper_title": "Sample Paper",
#         "chunk_id": 1,
#         "chunk_text": "This is a sample text",
#         "original_filename": "paper123.pdf"
#     }
# }]

# # 加载数据
# query_result = QueryResult(results=[
#     from_dict(data_class=QueryResultItem, data=item) 
#     for item in data
# ])
# print(query_result)
